'''
@author: zhangkai
@license: (C) Copyright 2017-2023
@contact: jeffcobile@gmail.com
@Software : PyCharm
@file: model_zoo.py
@time: 2020-05-29 11:19:38
@desc: 
'''
from ELib.utils.register import Registry
import torch
from collections import OrderedDict

MODEL_ZOO = Registry("MODEL")


def get_segmentation_model(cfg):
    """
    Built the whole model, defined by `cfg.MODEL.META_ARCHITECTURE`.
    """
    model_name = cfg.MODEL.NAME
    model = MODEL_ZOO.get(model_name)(cfg)
    # load_model_pretrain(cfg, model)
    return model


def load_model_pretrain(cfg, model):
    if cfg.BASE.PHASE == 'train':
        if cfg.TRAIN.PRETRAINED_MODEL_PATH:
            state_dict_to_load = torch.load(cfg.TRAIN.PRETRAINED_MODEL_PATH)
            keys_wrong_shape = []
            state_dict_suitable = OrderedDict()
            state_dict = model.state_dict()
            for k, v in state_dict_to_load.items():
                if v.shape == state_dict[k].shape:
                    state_dict_suitable[k] = v
                else:
                    keys_wrong_shape.append(k)
            model.load_state_dict(state_dict_suitable, strict=False)
    else:
        if cfg.TEST.TEST_MODEL_PATH:
            model.load_state_dict(torch.load(cfg.TEST.TEST_MODEL_PATH, map_location='cpu'), strict=False)